{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "6fa5cf6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import random\n",
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7d5f0bcc",
   "metadata": {},
   "outputs": [],
   "source": [
    " file = '../us_edges_id2newid.txt'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b289c07a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('us', 'idx2us.npy')"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#对边的编号\n",
    "type=file.split(\"/\")[1].split('_')[0]\n",
    "filename='idx2' + type + '.npy'\n",
    "type,filename"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "59680f78",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "11540\n",
      "打乱顺序前的数组: [    0     1     2 ... 11537 11538 11539]\n",
      "打乱顺序后的数组: [ 847 9474 1646 ...  173 1244 4989]\n",
      "6924 2308 2308\n",
      "train_idx： [    1     2     3 ... 11534 11535 11536]\n",
      "val_idx： [    0     7     9 ... 11537 11538 11539]\n",
      "test_idx： [    5    11    16 ... 11521 11523 11524]\n"
     ]
    }
   ],
   "source": [
    "dic={}\n",
    "with open(file,'r') as f:\n",
    "    reader=f.readlines()\n",
    "    edges_num=len(reader)\n",
    "    print(len(reader))\n",
    "    a=np.arange(edges_num)\n",
    "    print('打乱顺序前的数组:',a)\n",
    "    np.random.seed(2022)\n",
    "    np.random.shuffle(a)\n",
    "    print('打乱顺序后的数组:',a)\n",
    "    train_ratio,val_ratio,test_ratio=0.6,0.2,0.2\n",
    "    train_num=math.floor(edges_num*train_ratio)\n",
    "    val_num=math.floor(edges_num*val_ratio)\n",
    "    #print(train_num,val_num,edges_num-train_num-val_num)\n",
    "    train_idx = a[:train_num]\n",
    "    val_idx=a[train_num:(train_num+val_num)]\n",
    "    test_idx=a[(train_num+val_num):edges_num]\n",
    "    print(len(train_idx),len(val_idx),len(test_idx))\n",
    "    train_idx.sort()\n",
    "    val_idx.sort()\n",
    "    test_idx.sort()\n",
    "    print(\"train_idx：\",train_idx)\n",
    "    print(\"val_idx：\",val_idx)\n",
    "    print(\"test_idx：\",test_idx)\n",
    "    np.savez('train_val_test_idx.npz', train_idx=train_idx, val_idx=val_idx, test_idx=test_idx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "d41a1e2f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[    0     1     2 ... 11537 11538 11539]\n"
     ]
    }
   ],
   "source": [
    "new_train_idx=np.concatenate((train_idx,val_idx))\n",
    "new_train_idx.sort()\n",
    "print(new_train_idx)\n",
    "np.savez('train_test_idx.npz', train_idx=new_train_idx, test_idx=test_idx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "79810808",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "11053"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "d12b2c42",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "11053"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0ff0b39",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
